import numpy as np
from pymcr.mcr import McrAR
from input import input_data
from matplotlib import pyplot as plt
from pre.smoothing import savitzky_golay_filter
from pre.normalization import unit_vector_normalization
from pre.baseline_correction import polynomial_baseline_correction
from pymcr.regressors import OLS, NNLS
from pymcr.constraints import ConstraintNonneg, ConstraintNorm


def mcr_als(k:int,data:np.ndarray):
    '''
    MCR-ALS算法，输入数据集和可能的成分个数，预测每个成分的光谱与每个数据集中各个成分的占比
    :param k: int，预测的成分个数
    :param data: np.ndarray，n*m的矩阵，n为多少个数据，m为一个数据有多少个光强
    :return: s:np.ndarray 每个成分的光谱数据，k*m
    :return: c:np.ndarray 每个数据对应的各成分浓度 n*k
    '''
    mcr = McrAR(max_iter=100, st_regr='NNLS', c_regr=OLS(),
                c_constraints=[ConstraintNonneg(), ConstraintNorm()],
                tol_increase=1.0)

    n = data.shape[1]
    s_init = np.random.randint(0,100,size=(k,n))
    mcr.fit(data,ST=s_init)

    s = mcr.ST_opt_
    c = mcr.C_opt_
    row_sums = np.sum(c, axis=1)
    x = np.linspace(1, n, n)
    # plt.plot(x,data[0])
    # plt.show()
    #print(s[0])
    print(f"预测的各成分浓度分别为：{c}")
    #print(f"浓度总和：{row_sums}")
    return s,c


if __name__ == "__main__":
    data,_ = input_data()
    #pre
    #data = data[:,1000:]#切片去除激光部分
    #print(f"切片后的数据形状：{data.shape}")
    data = savitzky_golay_filter(data)#做SG平滑
    #print(data)
    data = unit_vector_normalization(data)#归一化
    #print(data)
    data = polynomial_baseline_correction(data)#基线矫正
    mcr_als(2,data)
